#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
文件类型统计脚本
用于统计test_dmp.csv文件中不同文件类型的数据量
"""
import csv
import os
from collections import Counter
import argparse


def count_file_types(csv_file_path):
    """
    统计CSV文件中不同文件类型的数据量
    
    Args:
        csv_file_path: CSV文件路径
        
    Returns:
        dict: 文件类型统计结果
    """
    file_types_count = Counter()
    
    # 检查文件是否存在
    if not os.path.exists(csv_file_path):
        print(f"错误: 文件 {csv_file_path} 不存在")
        return file_types_count
    
    try:
        # 读取CSV文件并统计文件类型
        with open(csv_file_path, 'r', encoding='utf-8', newline='') as csvfile:
            # 使用csv.reader读取文件，注意处理引号
            csv_reader = csv.reader(csvfile)
            # 跳过标题行
            next(csv_reader, None)
            
            for row in csv_reader:
                if len(row) >= 2:
                    file_path = row[1]  # 文件路径在第二列
                    # 获取文件扩展名
                    _, extension = os.path.splitext(file_path)
                    # 如果没有扩展名，使用'无扩展名'标记
                    if not extension:
                        extension = '无扩展名'
                    # 将扩展名转为小写并统计
                    file_types_count[extension.lower()] += 1
        
    except Exception as e:
        print(f"读取文件时发生错误: {e}")
    
    return file_types_count


def print_statistics(file_types_count, output_file=None):
    """
    打印或保存统计结果
    
    Args:
        file_types_count: 文件类型统计结果
        output_file: 输出文件路径，None则打印到控制台
    """
    total_count = sum(file_types_count.values())
    
    # 准备输出内容
    output_lines = ["文件类型统计结果", "="*50]
    output_lines.append(f"{'文件类型':<20}{'数量':<10}{'占比':<10}")
    output_lines.append("-"*50)
    
    # 按数量降序排列并添加到输出
    for file_type, count in sorted(file_types_count.items(), key=lambda x: x[1], reverse=True):
        percentage = (count / total_count) * 100 if total_count > 0 else 0
        output_lines.append(f"{file_type:<20}{count:<10}{percentage:.2f}%")
    
    output_lines.append("="*50)
    output_lines.append(f"{'总计':<20}{total_count:<10}{'100.00%':<10}")
    
    # 输出结果
    result = "\n".join(output_lines)
    
    if output_file:
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(result)
            print(f"统计结果已保存到: {output_file}")
        except Exception as e:
            print(f"保存文件时发生错误: {e}")
            print(result)
    else:
        print(result)
    
    return result


if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='统计CSV文件中不同文件类型的数据量')
    parser.add_argument('-f', '--file', default='test_dmp.csv', help='CSV文件路径，默认为test_dmp.csv')
    parser.add_argument('-o', '--output', help='输出文件路径，默认打印到控制台')
    args = parser.parse_args()
    
    # 获取CSV文件的绝对路径
    csv_file_path = os.path.abspath(args.file)
    
    print(f"正在统计文件: {csv_file_path}")
    
    # 统计文件类型
    file_types_count = count_file_types(csv_file_path)
    
    # 打印或保存统计结果
    print_statistics(file_types_count, args.output)